# 1.手动实现前馈神经网络解决多分类任务
from common import *


#  下载数据库
def down_load_data(path_dir):
    mnist_train = torchvision.datasets.MNIST(root=path_dir, train=True, download=True, transform=transforms.ToTensor())
    mnist_test = torchvision.datasets.MNIST(root=path_dir, train=False, transform=transforms.ToTensor())
    return mnist_train, mnist_test


#  读取小批量图集
def load_data_iter(mnist_train, mnist_test, batch_size, num_workers):
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return train_iter, test_iter


# 评估模型net在数据集data_iter中的准确率
def evaluate_accuracy(data_iter, net):
    loss_sum, acc_sum, n = 0.0, 0.0, 0
    for X, y in data_iter:
        y_hat = net(X)
        l = loss(y_hat, y).sum()
        loss_sum += l.item()
        acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
        n += y.shape[0]
    return acc_sum / n, loss_sum / n


def train(net, train_iter, test_iter, loss, num_epochs, batch_size,
              params=None, lr=None, optimizer=None):
    test_loss_avg = []
    train_loss_avg = []
    """ training network """
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            """ 第一步：计算模型输出和loss """
            y_hat = net(X)
            l = loss(y_hat, y).sum()

            # 梯度清零
            if optimizer is not None:
                optimizer.zero_grad()
            elif params is not None and params[0].grad is not None:
                for param in params:
                    param.grad.data.zero_()

            l.backward()  # 第二步 反向传播
            optimizer.step()

            train_l_sum += l.item()
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()  # 第三步 计算准确率
            n += y.shape[0]

        # 训练精确度
        test_acc, test_l = evaluate_accuracy(test_iter, net)
        train_loss_avg.append(train_l_sum / n)
        test_loss_avg.append(test_l)
        # 输出结果
        print('epoch %d, loss %.4f, train acc %.3f, test loss %.4f, test acc %.3f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_l, test_acc))
    return train_loss_avg, test_loss_avg


def get_fun(type):
    # 获取隐藏层
    if type == "sigmoid":
       return torch.nn.Sigmoid()
    elif type == "tanh":
        return torch.nn.Tanh()
    elif type == "leakyRelu":
        return torch.nn.LeakyReLU()
    elif type == "elu":
        return torch.nn.ELU()
    else:
        return torch.nn.ReLU()

num_inputs, num_hiddens, num_outs = 784, 10, 256
fun_t = "relu"

# 模型定义及参数初始化
net = torch.nn.Sequential(
    FlattenLayer(),
    torch.nn.Linear(num_inputs, num_hiddens),
    get_fun(fun_t),
    torch.nn.Linear(num_hiddens, num_outs)
)

for params in net.parameters():
    torch.nn.init.normal_(params, mean=0, std=0.01)
# 循环次数， 学习率
num_epochs, lr = 10, 0.1
# 交叉熵
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr)

# 创建测试数据及训练数据
batch_size = 32
num_workers = 0  # 多线程数
path_dir = '~/Datasets/MNIST'

mnist_train, mnist_test = down_load_data(path_dir)
train_iter, test_iter = load_data_iter(mnist_train, mnist_test, batch_size, num_workers)

# 开始训练 学习率0.1  训练次数5
train_loss, test_loss = train(net, train_iter, test_iter, loss, num_epochs, batch_size, params, lr, optimizer)
loss_draw(range(1, num_epochs + 1), train_loss, "epochs", "loss", range(1, num_epochs + 1), test_loss,
              ['train', 'test'])

